#ifndef Algebra_PetscSpace_h
#define Algebra_PetscSpace_h

#include "AlgebraModule.h"
#include "Space.h"

#include "Eigen/Dense"
#include "Eigen/Sparse"

// Forward declarations
typedef struct _p_Vec* Vec;
typedef struct _p_Mat* Mat;

SG_NAMESPACE_OPEN

namespace Algebra
{
    // Forward declarations
    template <typename TDataType> class PetscVector;
    template <typename TDataType> class PetscDenseMatrix;
    template <typename TDataType> class PetscSparseMatrix;

    /**
     * @brief Wrap PETSc parallel vector.
     *
     * @see https://petsc.org/release/manual/vec/
     *
     * @author nene
     * @date December, 2024.
     */
    template <typename TDataType>
    class ALGEBRA_EXPORT PetscVector : public Vector<TDataType>
    {
      public:
        typedef TDataType              DataType;
        typedef Vector<TDataType>      BaseType;
        typedef PetscVector<TDataType> ThisType;

        /**
         * @brief Proxy of one vector entry.
         *
         * Since access to PETSc vectors only goes through functions, rather than by obtaining
         * a reference to a vector element, we need a wrapper class that acts as if it was a
         * reference, and basically redirects all accesses (read and write) to member functions
         * of this class.
         *
         * @see https://github.com/dealii/dealii/blob/master/include/deal.II/lac/petsc_vector_base.h
         */
        class Entry
        {
            // Make the vector class a friend, so that it can create objects of the
            // present type.
            friend class PetscVector<TDataType>;

          private:
            /**
             * Constructor. It is made private so as to only allow the actual vector
             * class to create it.
             */
            Entry (const PetscVector<TDataType>& vector, const Index_t index) : mVector (vector), mIndex (index)
            {}

          public:
            /* Copy constructor. */
            Entry (const Entry& other) = default;

            // /**
            //  * This looks like a copy operator, but does something different than
            //  * usual. In particular, it does not copy the member variables of this
            //  * reference. Rather, it handles the situation where we have two vectors
            //  * @p v and @p w, and assign elements like in <tt>v(i)=w(i)</tt>. Here,
            //  * both left and right hand side of the assignment have data type
            //  * VectorReference, but what we really mean is to assign the vector
            //  * elements represented by the two references. This operator implements
            //  * this operation. Note also that this allows us to make the assignment
            //  * operator const.
            //  */
            // const Entry& operator= (const Entry& r) const;

            /**
             * The same function as above, but for non-const reference objects. The
             * function is needed since the compiler might otherwise automatically
             * generate a copy operator for non-const objects.
             */
            Entry& operator= (const Entry& other);

            /** Set the referenced element of the vector to <tt>s</tt>. */
            const Entry& operator= (const TDataType& s) const;

            /** Add <tt>s</tt> to the referenced element of the vector. */
            const Entry& operator+= (const TDataType& s) const;

            /** Subtract <tt>s</tt> from the referenced element of the vector. */
            const Entry& operator-= (const TDataType& s) const;

            /** Multiply the referenced element of the vector by <tt>s</tt>. */
            const Entry& operator*= (const TDataType& s) const;

            /** Divide the referenced element of the vector by <tt>s</tt>. */
            const Entry& operator/= (const TDataType& s) const;

            /**
             * Convert the reference to an actual value, i.e. return the value of
             * the referenced element of the vector.
             */
            operator TDataType () const;

            /** @name Maintain compatibilities with the std::complex<TDataType> data type. */
            //@{
            /** Return the real part of the value of the referenced element. */
            Real_t real () const;

            /** Return the imaginary part of the value of the referenced element. */
            Real_t imag () const;
            //@}

          private:
            const PetscVector<TDataType>& mVector;  //< Point to the vector we are referencing.
            const Index_t                 mIndex;   //< Index of the referenced element of the vector.
        };

        /// Default constructor
        PetscVector ();
        /// Construct a vector with given pattern.
        explicit PetscVector (Sparsity* sp);
        /// Construct a vector with given dimension.
        explicit PetscVector (const Index_t& n);
        /// Create holder of a PETSc Vec object/pointer. The Vec x object
        /// should already be created. If inc_ref_count is true, the reference
        /// counter of the Vec object will be increased. The Vec reference
        /// count will always be decreased upon destruction.
        PetscVector (Vec x, bool inc_ref_count);
        /// Copy constructor
        PetscVector (const PetscVector<TDataType>& x);
        /// Move constructor
        PetscVector (PetscVector<TDataType>&& x);

        ~PetscVector () override;

        // Assignment operator
        PetscVector<TDataType>& operator= (const PetscVector<TDataType>& x);

        /// Move Assignment
        PetscVector<TDataType>& operator= (PetscVector<TDataType>&& x);

        /* Conversion operator to the underlying PETSc vector. */
        operator const Vec& () const;
        operator Vec& ();

        /** Get the underlying raw PETSc Vec pointer. */
        Vec vec ();
        Vec vec () const;

        /** Create a copy of the vector. */
        PetscVector<TDataType> clone () const;

        /** The corrresponding linear space which this algebraic object pertained to.*/
        Space<TDataType>* getSpace () override;

        // /** Initialize this vector using the sparsity structure. */
        // void init (Sparsity* sp) override;

        /**
         * Calls the Vector's internal assembly routines, ensuring
         * that the values are consistent across processors.
         */
        void close () override;

        /**
         * Restores the \p Vector<T> to a pristine state.
         */
        void clear () override;

        /** Start assembling the elements stashed. */
        void beginAssembly () override;

        /** Finish assembling the elements stashed. */
        void endAssembly () override;

        /** Ensure that the values are consistent across processors. */
        void assemble () override;

        // /** Test whether the vector has been assembled. */
        // bool assembled () override;

        /** Resize *this to a vector of length size. */
        void resize (Index_t size) override;

        /** Resize the vector to size while retaining old values. */
        void conservativeResize (Index_t size) override;

        /** The number of rows of the vector */
        Index_t rows () const override;

        /** The index of the first vector element stored on this processor.*/
        Index_t row_start () const override;

        /** The index of the last vector element (+1) stored on this processor. */
        Index_t row_end () const override;

        /** Set all entries to 0. */
        void setZero () override;

        /** Set v(i) = \p value. */
        void set (const Index_t& i, const TDataType& value) override;

        /** Add \p value to the vector entry specified by \p i. */
        void add (const Index_t& i, const TDataType& value) override;

        /** The reference of the ith entry of the vector. */
        Entry operator() (const Index_t& i);

        /** A copy of the ith entry of the vector. */
        TDataType operator() (const Index_t& i) const override;

        /** The reference of the ith entry of the vector. */
        Entry operator[] (const Index_t& i);

        /** A copy of the ith entry of the vector.*/
        TDataType operator[] (const Index_t& i) const override;

        /** Print the elements of a vector to the given output stream. */
        void print (std::ostream& out) const override;

        /** Create a subvector using the indices in \p rows. */
        void create_subvector (Vector<TDataType>& subvector, const GeneralIndexSet& rows) const override;

        /** Create a subvector from the stride index set in \p rows. */
        void create_subvector (Vector<TDataType>& subvector, const StrideIndexSet& rows) const override;

        /** Determine the vector component with maximum real part \p value and its location \p index. */
        void max (Index_t& index, TDataType& value) const override;

        /**
         * The \f$ \ell_{\infty} \f$-norm of the vector, i.e. the maximum
         * absolute value of the entries of the vector.
         */
        void linfty_norm (TDataType& value) const override;

      private:
        struct PetscVectorPrivate* d = nullptr;
    };

    /**
     * @brief Wrap PETSc dense matrix.
     *
     * @see https://petsc.org/release/manual/mat/
     *
     * @author nene
     * @date December, 2024.
     */
    template <typename TDataType>
    class ALGEBRA_EXPORT PetscDenseMatrix : public DenseMatrix<TDataType>
    {
      public:
        typedef TDataType                   DataType;
        typedef DenseMatrix<TDataType>      BaseType;
        typedef PetscDenseMatrix<TDataType> ThisType;

        PetscDenseMatrix ();
        PetscDenseMatrix (Index_t rows, Index_t cols);
        virtual ~PetscDenseMatrix ();

        /** Print the elements of a matrix to the given output stream. */
        void print (std::ostream& out) const override;

      private:
        struct PetscDenseMatrixPrivate* d = nullptr;
    };

    /**
     * @brief Wrap PETSc parallel sparse matrix.
     *
     * The use of PETSc matrices involves the following actions: create a particular type of matrix,
     * insert values into it, process the matrix, use the matrix for various computations, and 
     * finally destroy the matrix.
     * 
     * @see https://petsc.org/release/manual/mat/
     *
     * @author nene
     * @date December, 2024.
     */
    template <typename TDataType>
    class ALGEBRA_EXPORT PetscSparseMatrix : public SparseMatrix<TDataType>
    {
      public:
        typedef TDataType                              DataType;
        typedef SparseMatrix<TDataType>                BaseType;
        typedef PetscSparseMatrix<TDataType>           ThisType;

        /**
         * @brief Proxy of one matrix element.
         */
        class Entry
        {
            friend class PetscSparseMatrix<TDataType>;

          private:
            Entry (PetscSparseMatrix<TDataType>& mat, int i, int j) : mMatrix (mat), mRow (i), mCol (j)
            {}

          public:
            /* Copy constructor. */
            Entry (const Entry& other) = default;

            /**
             * The same function as above, but for non-const reference objects. The
             * function is needed since the compiler might otherwise automatically
             * generate a copy operator for non-const objects.
             */
            Entry& operator= (const Entry& other);

            /** Set the referenced element of the matrix to <tt>s</tt>. */
            const Entry& operator= (const TDataType& s) const;

            /** Add <tt>s</tt> to the referenced element of the matrix. */
            const Entry& operator+= (const TDataType& s) const;

            /** Subtract <tt>s</tt> from the referenced element of the matrix. */
            const Entry& operator-= (const TDataType& s) const;

            /** Multiply the referenced element of the matrix by <tt>s</tt>. */
            const Entry& operator*= (const TDataType& s) const;

            /** Divide the referenced element of the matrix by <tt>s</tt>. */
            const Entry& operator/= (const TDataType& s) const;

            /**
             * Convert the reference to an actual value, i.e. return the value of
             * the referenced element of the matrix.
             */
            operator TDataType () const;

          private:
            const PetscSparseMatrix<TDataType>& mMatrix;  //< Point to the matrix we are referencing.
            const Index_t                       mRow;     //< row number.
            const Index_t                       mCol;     //< col number.
        };

        /// Default constructor
        PetscSparseMatrix ();
        /// Construct a matrix with given pattern.
        explicit PetscSparseMatrix (Sparsity* sp);
        /// Construct a vector with given dimension.
        PetscSparseMatrix (Index_t rows, Index_t cols);
        /// Create holder of a PETSc Mat object/pointer. The Mat A object
        /// should already be created. If inc_ref_count is true, the reference
        /// counter of the Mat will be increased. The Mat reference count will
        /// always be decreased upon destruction..
        PetscSparseMatrix (Mat A, bool inc_ref_count);
        /// Copy constructor
        PetscSparseMatrix (const PetscSparseMatrix<TDataType>& A);
        /// Move constructor
        PetscSparseMatrix (PetscSparseMatrix<TDataType>&& A);
        /// Destructor
        ~PetscSparseMatrix () override;

        /// Assignment operator
        PetscSparseMatrix<TDataType>& operator= (const PetscSparseMatrix<TDataType>& A);

        /// Move assignment operator
        PetscSparseMatrix<TDataType>& operator= (PetscSparseMatrix<TDataType>&& A);

        /** Create a copy of the vector. */
        PetscSparseMatrix<TDataType> clone () const;

        /** The corrresponding linear space which this algebraic object pertained to.*/
        Space<TDataType>* getSpace () override;

        /* Conversion operator to the underlying PETSc vector. */
        operator const Mat& () const;
        operator Mat& ();

        /** The underlying raw PETSc matrix pointer. */
        Mat mat ();
        Mat mat () const;

        // /** Initialize this matrix using the sparsity structure. */
        // void init (Sparsity* sp) override;

        /** Restore to a pristine state. */
        void clear () override;

        /**
         * Call the SparseMatrix's internal assembly routines, ensuring
         * that the values are consistent across processors.
         */
        void close () override;

        /** Start assembling the elements stashed. */
        void beginAssembly () override;

        /** Finish assembling the elements stashed. */
        void endAssembly () override;

        /** Ensure that the values are consistent across processors. */
        void assemble () override;

        /** Test whether the matrix has been assembled. */
        bool assembled () const override;

        /** Resize the matrix to a rows x cols matrix. */
        void resize (Index_t rows, Index_t cols) override;

        /** Resize the matrix to a rows x cols matrix leaving old values untouched. */
        void conservativeResize (Index_t rows, Index_t cols) override;

        /** The number of rows of the matrix */
        Index_t rows () const override;

        /** The number of columns of the matrix */
        Index_t cols () const override;

        /** The index of the first matrix row stored on this processor.*/
        Index_t row_start () const override;

        /** The index of the last matrix row (+1) stored on this processor. */
        Index_t row_end () const override;

        /** The index of the first matrix column owned by this processor. */
        Index_t col_start () const override;

        /** The index of the last matrix column (+1) owned by this processor.*/
        Index_t col_end () const override;

        /** Set all entries to 0. */
        void setZero () override;

        /** Set *this to the identity matrix. */
        void setIdentity () override;

        /**
         * Set the element \p (i,j) to \p value.  Throws an error if the
         * entry does not exist. Zero values can be "stored" in non-existent
         * fields.
         */
        void set (const Index_t& i, const Index_t& j, const TDataType& value) override;

        /**
         * Add \p value to the element \p (i,j).  Throws an error if the
         * entry does not exist. Zero values can be "added" to non-existent
         * entries.
         */
        void add (const Index_t& i, const Index_t& j, const TDataType& value) override;

        // /**
        //  * Add the full matrix \p dm to the SparseMatrix.  This is useful
        //  * for adding an element matrix at assembly time.
        //  */
        // void add (const DenseMatrix<TDataType>& dm, const IndexMap_t& rows, const IndexMap_t& cols) override;

        // /**
        //  * Add \p spm to the rows and cols of this matrix (A)
        //  */
        // void add (const SparseMatrix<TDataType>& /*spm*/, const IndexMap_t& /*rows*/, const IndexMap_t& /*cols*/) override;

        /** The reference of matrix entry \p(i,j). */
        Entry operator() (const Index_t& i, const Index_t& j);

        /**
         * \returns A copy of matrix entry \p (i,j).
         *
         * \note This may be an expensive operation, and you should always
         * be careful where you call this function.
         */
        TDataType operator() (const Index_t& i, const Index_t& j) const override;

        /** The number of nonzeros. */
        Index_t nonZeros () const override;

        /** Get the CSR (Compressed Sparse Row) Storage of the sparse matrix. */
        void get_csr (std::vector<Index_t>& ia, std::vector<Index_t>& ja, std::vector<TDataType>& a) const override;

        /** Print the elements of a matrix to the given output stream. */
        void print (std::ostream& out) const override;

        /** Create a submatrix from the row/col general-index-set. */
        void create_submatrix (SparseMatrix<TDataType>& submatrix, const GeneralIndexSet& rows, const GeneralIndexSet& cols) const override;

        /** Create a submatrix from the row/col stride-index-set. */
        void create_submatrix (SparseMatrix<TDataType>& submatrix, const StrideIndexSet& rows, const StrideIndexSet& cols) const override;

        /**
         * Similar to the above function, this function creates a \p
         * submatrix which is defined by the indices given in the \p rows
         * and \p cols vectors.
         * Note: Both \p rows and \p cols can be unsorted;
         *       Use the above function for better efficiency if your indices are sorted;
         *       \p rows and \p cols can contain indices that are owned by other processors.
         */
        void create_submatrix_nosort (SparseMatrix<TDataType>& /*submatrix*/,
                                      const GeneralIndexSet& /*rows*/,
                                      const GeneralIndexSet& /*cols*/) const override;

        /**
         * Create a submatrix from the row/col stride-index-set.
         * \p rows and \p cols can contain indices that are owned by other processors.
         *
         */
        void create_submatrix_nosort (SparseMatrix<TDataType>& /*submatrix*/,
                                      const StrideIndexSet& /*rows*/,
                                      const StrideIndexSet& /*cols*/) const override;

        /** Copy the diagonal part of the matrix into \p dest. */
        void get_diagonal (Vector<TDataType>& dest) const override;

        /** Copy the transpose of the matrix into \p dest, which may be *this. */
        void get_transpose (SparseMatrix<TDataType>& dest) const override;

        /** Copy the inverse of this matrix into \p dest, which may be *this. */
        void get_inverse (SparseMatrix<TDataType>& dest) const override;

        /** Scale all elements of this matrix by \p scale */
        void scale (const TDataType scale) override;

        /** Compute Y = A*X for matrix \p X. */
        void matrix_matrix_mult (SparseMatrix<TDataType>& X, SparseMatrix<TDataType>& Y) override;

      private:
        struct PetscSparseMatrixPrivate* d = nullptr;
    };

    template <typename TDataType>
    PetscVector<TDataType> operator+ (const PetscVector<TDataType>&, const PetscVector<TDataType>&);
    template <typename TDataType>
    PetscVector<TDataType> operator- (const PetscVector<TDataType>&, const PetscVector<TDataType>&);
    template <typename TDataType>
    PetscVector<TDataType> operator* (const TDataType&, const PetscVector<TDataType>&);
    template <typename TDataType>
    TDataType operator* (const PetscVector<TDataType>&, const PetscVector<TDataType>&);
    template <typename TDataType>
    PetscVector<TDataType> operator/ (const PetscVector<TDataType>&, const TDataType&);

    template <typename TDataType>
    PetscSparseMatrix<TDataType> operator+ (const PetscSparseMatrix<TDataType>&, const PetscSparseMatrix<TDataType>&);
    template <typename TDataType>
    PetscSparseMatrix<TDataType> operator- (const PetscSparseMatrix<TDataType>&, const PetscSparseMatrix<TDataType>&);
    template <typename TDataType>
    PetscSparseMatrix<TDataType> operator* (const TDataType&, const PetscSparseMatrix<TDataType>&);
    template <typename TDataType>
    PetscVector<TDataType> operator* (const PetscSparseMatrix<TDataType>&, const PetscVector<TDataType>&);

    /** Compute the matrix quadratic value \f $\boldsymbol{x}^{T}\boldsymbol{A}\boldsymbol{x} \f$ */
    template <typename TDataType>
    TDataType quadratic (const PetscSparseMatrix<TDataType>& A, const PetscVector<TDataType>& x);

    template <typename TDataType>
    std::ostream& operator<< (std::ostream& s, const PetscVector<TDataType>& v);
    template <typename TDataType>
    std::ostream& operator<< (std::ostream& s, const PetscDenseMatrix<TDataType>& m);
    template <typename TDataType>
    std::ostream& operator<< (std::ostream& s, const PetscSparseMatrix<TDataType>& m);

    /** Pack the wrappers of PETSc algebraic objects into a space. */
    template <typename TDataType>
    class ALGEBRA_EXPORT PetscSpace : public Space<TDataType>
    {
      public:
        using DataType         = TDataType;
        using VectorType       = PetscVector<TDataType>;
        using DenseMatrixType  = PetscDenseMatrix<TDataType>;
        using SparseMatrixType = PetscSparseMatrix<TDataType>;
    };

    // // Enforce the explicit instantiations to hide the underlying implementation.
    // template class ALGEBRA_EXPORT PetscVector<Real_t>;
    // template class ALGEBRA_EXPORT PetscVector<Complex_t>;

    // template class ALGEBRA_EXPORT PetscDenseMatrix<Real_t>;
    // template class ALGEBRA_EXPORT PetscDenseMatrix<Complex_t>;

    // template class ALGEBRA_EXPORT PetscSparseMatrix<Real_t>;
    // template class ALGEBRA_EXPORT PetscSparseMatrix<Complex_t>;

    // template class ALGEBRA_EXPORT PetscSpace<Real_t>;
    // template class ALGEBRA_EXPORT PetscSpace<Complex_t>;

    // Export the data types for external modules.
    template <typename TDataType>
    using Vector_t = typename PetscSpace<TDataType>::VectorType;

    template <typename TDataType>
    using Matrix_t = typename PetscSpace<TDataType>::DenseMatrixType;

    template <typename TDataType>
    using SparseMatrix_t = typename PetscSpace<TDataType>::SparseMatrixType;

    // real types
    typedef Vector_t<Real_t>       Vec_t;
    typedef Matrix_t<Real_t>       Mat_t;
    typedef SparseMatrix_t<Real_t> SpMat_t;

    // complex types
    typedef Vector_t<Complex_t>       ComplexVec_t;
    typedef Matrix_t<Complex_t>       ComplexMat_t;
    typedef SparseMatrix_t<Complex_t> ComplexSpMat_t;
}  // namespace Algebra

SG_NAMESPACE_CLOSE

#endif  // Algebra_PetscSpace_h